import asyncio
import websockets
import json


async def handle_connection(websocket, path):
    """Handles incoming websocket connections."""

    print(f"New connection from: {websocket.remote_address}")

    # Access headers from handshake request
    headers = websocket.request_headers
    print(f"Received headers: {dict(headers)}")

    try:
        async for message in websocket:
            print(f"Received message: {message}")
            response = f"Server received: {message} with headers: {dict(headers)}"
            await websocket.send(response)

    except websockets.exceptions.ConnectionClosed as e:
        print(f"Connection closed. Code: {e.code}, Reason: {e.reason}")


async def main():
    """Starts the WebSocket server."""
    async with websockets.serve(handle_connection, "localhost", 8765) as ws_server:
        print("WebSocket server started at ws://localhost:8765")
        await ws_server.wait_closed()


if __name__ == "__main__":
    asyncio.run(main())
